import tensorflow as tf
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data

"""
    CNN卷积神经网络：
        conv2d(卷积) -> relu(激活函数) -> pool(池化)        多次重复，将数据样本的shape减少。
        
        
        处理数据shape：
        
        matmul(pool,weight)  + basic  
        
        
        loss:
        
        train:
        
        
         

    激活函数
    Sigmoid/Tanh/ReLU
    
    激活函数(Activation Function)的特点：
        非线性： 当激活函数是线性的时候，一个两层的神经网络就可以逼近基本上所有的函数了。
        可微： 当优化方法是基于梯度的时候，这个性质是必须的。
        单调性： 当激活函数是单调的时候，单层网络能够保证是凸函数。
        输出值范围： 当激活函数输出值是 有限 的时候，基于梯度的优化方法会更加 稳定，因为特征的表示受有限权值的影响更显著；当激活函数的输出是 无限 的时候，模型的训练会更加高效，不过在这种情况小，一般需要更小的学习率。
    
"""

"""
    初始化权值
"""


def weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev=0.1)  # 生成一个截断的正态分布
    return tf.Variable(initial)


"""
    初始化偏置
"""


def bias_variable(shape):
    initial = tf.constant(0.1, shape=shape)
    return tf.Variable(initial)


"""
    用比较老的方式训练
"""


def train_old_way():
    mnist = input_data.read_data_sets('MNIST_data', one_hot=True)

    # 输入、输入数据处理
    x_data = tf.placeholder(tf.float32, [None, 28 * 28])
    y_data = tf.placeholder(tf.float32, [None, 10])
    # 改变x的格式转为4D的向量[batch, in_height, in_width, in_channels]`
    x_image = tf.reshape(x_data, [-1, 28, 28, 1])

    """"卷积操作:
        每张图片生成32个特征模型，因为每个图片的卷积核为32个。 32 * 28 * 28
    """
    # `strides[0] = strides[3] = 1`. strides[1]代表x方向的步长，strides[2]代表y方向的步长
    # 利用32个卷积核计算，
    # 每个卷积核一个偏置值,每一个通道一个偏置值
    # weight_variable([5, 5, 1, 32]) ：5*5的采样窗口，32个卷积核从1个平面抽取特征
    relu_conv2d = tf.nn.relu(
        tf.nn.conv2d(x_image, weight_variable([5, 5, 1, 32]), strides=[1, 1, 1, 1], padding="SAME") + bias_variable(
            [32]))

    """池化操作：
       特征模型数目不变，但大小减小为原来的一般，因为步长为2。  32 * 14 *14
       `strides[0] = strides[3] = 1`. strides[1]代表x方向的步长，strides[2]代表y方向的步长.[中间代表x,y]
    """
    pool_1 = tf.nn.max_pool(relu_conv2d, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")

    """"
        第二次卷积：生成64个特征模型
    """
    conv2d_2 = tf.nn.conv2d(pool_1, weight_variable([5, 5, 32, 64]), strides=[1, 1, 1, 1],
                            padding="SAME") + bias_variable([64])

    relu_conv2d_2 = tf.nn.relu(conv2d_2)

    pool_2 = tf.nn.max_pool(relu_conv2d_2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")

    """
        将样本数据转化成和y_data 一样的数据shape。
    """
    # tf.reshape(pool_2, [-1, 64 * 7 * 7] 将池化后的每个样本数据压缩成一维数据
    h_fc1 = tf.matmul(tf.reshape(pool_2, [-1, 64 * 7 * 7]), weight_variable([7 * 7 * 64, 1024])) + bias_variable([1024])
    h_fc1_dropout = tf.nn.dropout(tf.nn.relu(h_fc1), 1.0)
    h_fc2 = tf.matmul(h_fc1_dropout, weight_variable([1024, 10])) + bias_variable([10])

    # 计算出0-1的 softmax 值，得到所有概率的候选值。e的x次方 / e次方[xi到xn]和
    prediction = tf.nn.softmax(h_fc2)

    # 损失值计算、训练定义
    cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y_data, logits=prediction))
    train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy)

    # 准确率  argmax返回一维张量中最大的值所在的位置，即1所在的位置
    correct_prediction = tf.equal(tf.argmax(prediction, 1), tf.argmax(y_data, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

    with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as session:
        session.run(tf.global_variables_initializer())

        for i in range(1, 2001):
            x_input, y_input = mnist.train.next_batch(100)
            session.run(train_step, {x_data: x_input, y_data: y_input})
            if i % 100 == 0:
                test_acc = session.run(accuracy, feed_dict={x_data: mnist.test.images, y_data: mnist.test.labels})
                print("Iter " + str(i) + ", Testing Accuracy= " + str(test_acc))


"""
    用1.6之后的方法训练
"""


def train_new_way():
    train_data = tf.placeholder(tf.float32, [None, 28 * 28])
    train_labels = tf.placeholder(tf.float32, [None, 10])

    train_data_images = tf.reshape(train_data, [-1, 28, 28, 1])

    # [-1,28,28,32]
    conv2d_1 = tf.layers.conv2d(train_data_images, filters=32, kernel_size=[5, 5], strides=(1, 1), padding="SAME",
                                activation=tf.nn.relu)

    # [-1,14,14,32]
    max_pooling2d_1 = tf.layers.max_pooling2d(conv2d_1, pool_size=[2, 2], strides=2, padding="SAME")

    # [-1,14,14,64]
    conv2d_2 = tf.layers.conv2d(max_pooling2d_1, filters=64, kernel_size=[5, 5], padding="SAME", activation=tf.nn.relu)

    # [-1,7,7,64]
    max_pooling2d_2 = tf.layers.max_pooling2d(conv2d_2, pool_size=[2, 2], strides=2)

    # 将数据shape转化二维
    pool2_flat = tf.reshape(max_pooling2d_2, shape=[-1, 7 * 7 * 64])

    # [-1, 1024]
    dense_1024 = tf.layers.dense(pool2_flat, units=1024, activation=tf.nn.relu)

    # [-1, 10]
    prediction = tf.nn.softmax(tf.layers.dense(dense_1024, units=10))

    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=train_labels, logits=prediction))
    train_step = tf.train.AdamOptimizer(1e-4).minimize(loss)

    # 准确率[-1, 1]
    correct_prediction = tf.equal(tf.argmax(prediction, 1), tf.argmax(train_labels, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

    with tf.Session() as  session:
        session.run(tf.global_variables_initializer())
        mnist = input_data.read_data_sets('MNIST_data', one_hot=True)
        for step in range(1, 1001):
            x_data, y_data = mnist.train.next_batch(100)
            session.run(train_step, feed_dict={train_data: x_data, train_labels: y_data})
            if step % 100 == 0:
                accuracy_vale = session.run(accuracy,
                                            feed_dict={train_data: mnist.test.images, train_labels: mnist.test.labels})
                print("step=", step, "\taccuracy_vale=", accuracy_vale)


if __name__ == '__main__':
    train_new_way()
